import os
import random
from tqdm import tqdm
import SimpleITK as sitk
import cv2
import numpy as np


def SplitDataset(img_path, train_percent=0.9):
    data = os.listdir(img_path)
    train_images = []
    test_images = []
    num = len(data)
    train_num = int(num * train_percent)
    indexes = list(range(num))
    train = random.sample(indexes, train_num)
    for i in indexes:
        if i in train:
            train_images.append(data[i])
        else:
            test_images.append(data[i])
    return train_images, test_images


def conver(img_path, save_dir, mask_path=None, select_condition=None, mode="trian"):
    os.makedirs(save_dir, exist_ok=True)
    if mode == "train":
        savepath_img = os.path.join(save_dir, 'imagesTr')
        savepath_mask = os.path.join(save_dir, 'labelsTr')
    elif mode == "test":
        savepath_img = os.path.join(save_dir, 'imagesTs')
        savepath_mask = os.path.join(save_dir, 'labelsTs')
    os.makedirs(savepath_img, exist_ok=True)
    if mask_path is not None:
        os.makedirs(savepath_mask, exist_ok=True)

    ImgList = os.listdir(img_path)
    with tqdm(ImgList, desc="conver") as pbar:
        for name in pbar:
            if select_condition is not None and name not in select_condition:
                continue
            Img = cv2.imread(os.path.join(img_path, name))
            if mask_path is not None:
                Mask = cv2.imread(os.path.join(mask_path, name), 0)
                Mask = (Mask / 255).astype(np.uint8)
                if Img.shape[:2] != Mask.shape:
                    Mask = cv2.resize(Mask, (Img.shape[1], Img.shape[0]))
            Img_Transposed = np.transpose(Img, (2, 0, 1))
            Img_0 = Img_Transposed[0].reshape(1, Img_Transposed[0].shape[0], Img_Transposed[0].shape[1])
            Img_1 = Img_Transposed[1].reshape(1, Img_Transposed[1].shape[0], Img_Transposed[1].shape[1])
            Img_2 = Img_Transposed[2].reshape(1, Img_Transposed[2].shape[0], Img_Transposed[2].shape[1])
            if mask_path is not None:
                Mask = Mask.reshape(1, Mask.shape[0], Mask.shape[1])

            Img_0_name = name.split('.')[0] + '_0000.nii.gz'
            Img_1_name = name.split('.')[0] + '_0001.nii.gz'
            Img_2_name = name.split('.')[0] + '_0002.nii.gz'
            if mask_path is not None:
                Mask_name = name.split('.')[0] + '.nii.gz'

            Img_0_nii = sitk.GetImageFromArray(Img_0)
            Img_1_nii = sitk.GetImageFromArray(Img_1)
            Img_2_nii = sitk.GetImageFromArray(Img_2)
            if mask_path is not None:
                Mask_nii = sitk.GetImageFromArray(Mask)

            sitk.WriteImage(Img_0_nii, os.path.join(savepath_img, Img_0_name))
            sitk.WriteImage(Img_1_nii, os.path.join(savepath_img, Img_1_name))
            sitk.WriteImage(Img_2_nii, os.path.join(savepath_img, Img_2_name))
            if mask_path is not None:
                sitk.WriteImage(Mask_nii, os.path.join(savepath_mask, Mask_name))


if __name__ == "__main__":
    train_percent = 0.9
    img_path = r".../img"
    mask_path = r".../mask"
    output_folder = r"./dataset"
    os.makedirs(output_folder, exist_ok=True)
    train_images, test_images = SplitDataset(img_path, train_percent)
    conver(img_path, output_folder, mask_path, train_images, mode="train")
    conver(img_path, output_folder, mask_path, test_images, mode="test")
